from torch import nn

class MobileNetV2(nn.Module):
    # 用于控制总参数量
    alpha = 1
    beta = 1
    def __init__(self, classes):
        super().__init__()
        self.classes = classes
        self.feature, self.avg, self.fc = self.__create_nn()
        self.__init_args()

    def __create_nn(self):
        layer = nn.Sequential()
        layer.add_module("conv1", self._Conv2DBlock(3, int(32 * self.alpha), 3, 2, 1))
        layer.add_module("bottleneck_1", self._make_block(Bottleneck, [32, 16, [1], 1]))
        layer.add_module("bottleneck_2", self._make_block(Bottleneck, [16, 24, [2, 1], 6]))
        layer.add_module("bottleneck_3", self._make_block(Bottleneck, [24, 32, [2, 1, 1], 6]))
        layer.add_module("bottleneck_4", self._make_block(Bottleneck, [32, 64, [2, 1, 1, 1], 6]))
        layer.add_module("bottleneck_5", self._make_block(Bottleneck, [64, 96, [1, 1, 1], 6]))
        layer.add_module("bottleneck_6", self._make_block(Bottleneck, [96, 160, [2, 1, 1], 6]))
        layer.add_module("bottleneck_7", self._make_block(Bottleneck, [160, 320, [1], 6]))
        layer.add_module("conv2", self._Conv2DBlock(320, int(1280 * self.alpha), 1, 1))
        avg_pooling = nn.AdaptiveAvgPool2d((1, 1))
        fc = nn.Linear(int(1280 * self.alpha), self.classes)
        return layer, avg_pooling, fc

    def _Conv2DBlock(self, in_c, out_c, k, s, p=0, groups=1):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, k, s, p, groups=groups, bias=False),
            nn.BatchNorm2d(out_c),
            nn.ReLU6(True)
        )

    def _make_block(self, block, args):
        in_c, out_c, strides, t = tuple(args)
        layer = []
        for stride in strides:
            layer.append(block(in_c, out_c, stride, t))
            in_c = out_c

        return nn.Sequential(*layer)

    def __init_args(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=1e-3)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.feature(x)
        out = self.avg(out).view(out.size(0), -1)
        out = self.fc(out)
        return out

class Bottleneck(nn.Module):
    def __init__(self, c1, c2, stride, t):
        super().__init__()
        self.c1, self.c2, self.stride = c1, c2, stride
        self.m = nn.Sequential(
            nn.Conv2d(c1, c2 * t, 1, 1, bias=False),
            nn.BatchNorm2d(c2 * t),
            nn.ReLU6(True),
            nn.Conv2d(c2 * t, c2 * t, 3, stride, 1, groups=c2 * t, bias=False),
            nn.BatchNorm2d(c2 * t),
            nn.ReLU6(True),
            nn.Conv2d(c2 * t, c2, 1, 1, bias=False),
        )

        self.relu6 = nn.ReLU6(True)

    def forward(self, x):
        identify = x
        if self.stride == 1 and self.c1 == self.c2:
            out = self.m(x) + identify
        else:
            out = self.m(x)
        return out


def mobilenetv2(classes):
    return MobileNetV2(classes)

if __name__ == "__main__":
    from utils.utils import get_model_flops_args
    model = MobileNetV2(5)
    get_model_flops_args(model, (3, 224, 224))
